import torch
import torchvision.transforms as transforms
import numpy as np
import cv2
import logging

from .model import Net

# torchvision.transforms主要是用于常见的一些图形变换
# 模型训练是按照传统ReID的方法进行
# 使用Extractor类的时候输入为一个list的图片
# 得到图片对应的特征。

class Extractor(object):
    
    def __init__(self, model_path, use_cuda=True):
        # 初始化一个ReID网络
        self.net = Net(reid=True)
        # 指定使用cpu还是gpu
        self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
        # 模型加载-此处加载的应该是ReID的权重参数
        state_dict = torch.load(model_path, map_location=torch.device(self.device))[
            'net_dict']
        self.net.load_state_dict(state_dict)
        # 日志初始化
        logger = logging.getLogger("root.tracker")
        logger.info("Loading weights from {}... Done!".format(model_path))
        self.net.to(self.device) # 这里将网络转移到device上,准备使用gpu/cpu计算
        self.size = (64, 128) # 图像统一尺寸
        self.norm = transforms.Compose([
            transforms.ToTensor(), # 将PIL或NDarray转换为Tensor
            # 图像归一化
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

    def _preprocess(self, im_crops):
        """
        预处理,主要是对输入的im_crops处理

        TODO:
            1. to float with scale from 0 to 1
            2. resize to (64, 128) as Market1501 dataset did
            3. concatenate to a numpy array
            3. to torch Tensor
            4. normalize
        """
        def _resize(im, size):
            return cv2.resize(im.astype(np.float32)/255., size)

        im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(
            0) for im in im_crops], dim=0).float()
        return im_batch

    def __call__(self, im_crops):
        # 直接调用该对象时,可以看下方__name__=='__main__'的示例
        # 启用模型
        im_batch = self._preprocess(im_crops)
        with torch.no_grad():
            # 将数据搬移到device上
            im_batch = im_batch.to(self.device)
            # 在device上运算得到特征embedding
            features = self.net(im_batch)
        # 再将数据从device上拷贝到cpu上并转换成numpy格式输出
        return features.cpu().numpy()


if __name__ == '__main__':
    img = cv2.imread("demo.jpg")[:, :, (2, 1, 0)]
    extr = Extractor("checkpoint/ckpt.t7")
    feature = extr(img)
    print(feature.shape)
